#!/usr/bin/perl

# ------------------------------------------------------------------------------
# --- Finds pdb file's interchain partners and adds them to database -----------
# --- Original by Andrea 01/06/09 ----------------------------------------------
# --- Re-written by C. Bergey; documentation added by A. Watkins. --------------
# ------------------------------------------------------------------------------

use strict;
use warnings;

use DBI;

# Get PDB ID

my $pdb_code = shift;		# e.g. 3SL9, 1YCR, ...
chomp $pdb_code;

defined ($pdb_code) or 
	die "ERROR: Must pass PDB code as input.\n";

print STDERR "Finding interchain partners for PDB file $pdb_code...\n";

# e.g. pdb_files/3SL9.pdb, pdb_files/1YCR.pdb, ...
my $pdb_filename = "pdb_files/" . $pdb_code . ".pdb";

# Open each current PDB file

open (PDB, "<$pdb_filename")
	or die "ERROR: Could not open PDB file, $pdb_filename.\n";
		
# Array of arrays to hold the associated chains in the PDB file
my @compnd_chains;

foreach my $pdb_line (<PDB>) {
	
	chomp $pdb_line;
		
	# Look for lines that look like this:
	# COMPND   3 CHAIN: A, D;
	if ($pdb_line =~ /^COMPND.{5}CHAIN:/) {
			
		# Isolate chain names (one alphanumeric character)
			
		# First, take the portion of the line bit after "CHAIN:"
		my @chain_string = split (/:/, $pdb_line);
			
		# Then, we can acquire each code; put them in @chains
		my @chains = $chain_string[1] =~ /\s([\d\w])[,;]/g;
		
		# Add this line's chains to the big compound chain array for this PDB
		push (@compnd_chains, \@chains);
	}
}

# Now that we've all the possible interface chains, we can infer all
# the possible ways they could interact. Taking 1A0H as an example, the
# COMPND CHAIN lines are:

#	COMPND   3 CHAIN: A, D;                                                         
#	COMPND   9 CHAIN: B, E;                                                         

# So eventually we'll write the following lines to interface_chains_pdb.txt:
#  pdb_partner_files/1A0H_A_B.pdb		pdb_partner_files/1A0H_B_D.pdb
#  pdb_partner_files/1A0H_A_D.pdb		pdb_partner_files/1A0H_B_E.pdb
#  pdb_partner_files/1A0H_A_E.pdb		pdb_partner_files/1A0H_D_E.pdb

# Make an array containing all one-character chain IDs for this PDB
	
my @all_chains = ();
	
foreach my $chain_ref (@compnd_chains) {
	
	my @these_chains = @$chain_ref;
	@all_chains = (@all_chains, @these_chains);
}
	
# Sort @all_chains, so we don't look for B_A instead of A_B
@all_chains = sort(@all_chains);

# Now we need to update the backend database with the number of chains and their codes

my $dbh = DBI->connect(
	"dbi:SQLite:dbname=workflow_tracker.db",
	"",
	"",
	{ RaiseError => 1 },
) or die $DBI::errstr;

# First, number of chains

my $num_chains_sql = "UPDATE pdb SET num_chains = " . scalar (@all_chains);
$num_chains_sql .= " WHERE pdb.pdb_id = '" . $pdb_code . "';";

print STDERR $num_chains_sql . "\n\n";

my $sth = $dbh->prepare($num_chains_sql);
$sth->execute();
$sth->finish();

# Next

my $chains_sql = "UPDATE pdb SET chains = \"" . join (',', @all_chains) . "\"";
$chains_sql .= " WHERE pdb.pdb_id = '" . $pdb_code . "';";

print STDERR $chains_sql . "\n\n";

$sth = $dbh->prepare($chains_sql);
$sth->execute();
$sth->finish();

# Add all possible chain combos to the database
for (my $partner_1 = 0; $partner_1 < scalar (@all_chains); $partner_1++) {
	for (my $partner_2 = $partner_1 + 1; $partner_2 < scalar (@all_chains); $partner_2++) {
			
		# Figure out ID of this PDB in the database
		my $pdb_id_sql = "SELECT * FROM pdb WHERE pdb_id = '$pdb_code';";
		
		my $sth = $dbh->prepare($pdb_id_sql);
		$sth->execute();

		my @pdb_id_array = $sth->fetchrow_array();
		my $pdb_id = $pdb_id_array[0];
		
		$sth->finish();
				
		# Insert into PDB partner table of database
		my $chains_sql = "INSERT INTO pdb_partner (id, pdb_id, chain_1, chain_2, rosetta_status, ";
		$chains_sql .= "rosetta_job_id, rosetta_attempts, proc_ala_status, proc_ala_job_id, ";
		$chains_sql .= "proc_ala_attempts) ";
		$chains_sql .= "VALUES (NULL, \"$pdb_id\", \"$all_chains[$partner_1]\", \"$all_chains[$partner_2]\", \"\", ";
		$chains_sql .= "NULL, 0, \"\", NULL, 0);\n";
		
		my $sth = $dbh->prepare($chains_sql);
		$sth->execute();

		$sth->finish();

	}
}
		
close PDB;

$dbh->disconnect();


exit;
